package flowlogic;

import java.util.ArrayList;
import java.util.HashMap;
import java.util.HashSet;

import util.Pair;

import ast.ActsFor;
import ast.AndPredicate;
import ast.Assert;
import ast.AtomicLabel;
import ast.Binder;
import ast.Call;
import ast.Case;
import ast.ChannelConstant;
import ast.CompositeBinder;
import ast.Constant;
import ast.DQCSystem;
import ast.DataConstant;
import ast.Dummy;
import ast.Endorse;
import ast.ExpVar;
import ast.Expression;
import ast.FunApp;
import ast.Function;
import ast.Go;
import ast.In;
import ast.Label;
import ast.LabelConjunction;
import ast.LabelDisjunction;
import ast.Network;
import ast.NewConstN;
import ast.NewConstP;
import ast.None;
import ast.NumPredicate;
import ast.OrPredicate;
import ast.Out;
import ast.ParallelNetwork;
import ast.ParallelProcess;
import ast.PrefixedProcess;
import ast.Principal;
import ast.Procedure;
import ast.ProcessAtSite;
import ast.Process;
import ast.QualityPredicate;
import ast.Site;
import ast.Some;
import ast.Term;
import ast.TermVar;
import ast.Zero;

public class STSimplifier {
	
	private static HashMap<String, ChannelConstant> channelMap = 
			new HashMap<String, ChannelConstant>();
	private static HashMap<String, DataConstant> dataMap = 
			new HashMap<String, DataConstant>();
	private static HashSet<String> principals = new HashSet<String>();
	private static HashSet<String> sites = new HashSet<String>();
	
	private static HashMap<String, Function> functionMap =
			new HashMap<String, Function>();
	
	private static HashMap<String, Procedure> procedureMap = 
			new HashMap<String, Procedure>();
	
	public static HashMap<String, ChannelConstant> getChannelMap()
	{
		return channelMap;
	}
	
	public static HashMap<String, DataConstant> getDataMap()
	{
		return dataMap;
	}
	
	public static HashSet<String> getPrincipals()
	{
		return principals;
	}
	
	public static HashSet<String> getSites()
	{
		return sites;
	}
	
	public static HashMap<String, Function> getFunctionMap() {
		return functionMap;
	}
	
	public static HashMap<String, Procedure> getProcedureMap() {
		return procedureMap;
	}

	public static String getChannelMapString()
	{
		String res = "";
		for (String str : channelMap.keySet()) {
			res += str + " -> ";
			ChannelConstant cst = channelMap.get(str);
			String ext = (cst.getLabel().isExternal() ? "ext" : "int");
			res += cst.getName() + "(" + ext + "):" + cst.getLabel() + "\n";
		}
		return res;
	}
	
	public static String getDataMapString()
	{
		String res = "";
		for (String str : dataMap.keySet()) {
			res += str + " ";
		}
		return res;
	}
	
	public static String getPrincipalsString()
	{
		String res = "";
		for (String str : principals) {
			res += str + " ";
		}
		return res;
	}
	
	public static String getSitesString()
	{
		String res = "";
		for (String str : sites) {
			res += str + " ";
		}
		return res;
	}
	
	public static void traverseSystem(DQCSystem sys)
	{
		
		channelMap.clear();
		dataMap.clear();
		principals.clear();
		sites.clear();
		functionMap.clear();
		
		ArrayList<Constant> clsys = sys.getClsys();
		for (Constant cst : clsys) {
			if (cst instanceof ChannelConstant) {
				Label lb = ((ChannelConstant) cst).getLabel();
				Label simplified = simplifyLabel(lb);
				simplified.setExternal(lb.isExternal());
				((ChannelConstant) cst).setLabel(simplified);
				
				channelMap.put(cst.getName(), (ChannelConstant) cst);
			} else {
				dataMap.put(cst.getName(), (DataConstant) cst);
			}
		}
		
		for (Procedure proc : sys.getProcedures()) {
			proc.setProcess(simplifyDegenerateParProcess(proc.getProcess()));
			traverseProcess(proc.getProcess());
			
			sites.add(proc.getSite().getName());
			principals.add(proc.getSite().getName());
			
			proc.setBeginLabel(simplifyLabel(proc.getBeginLabel()));
		}
		
		traverseNetwork(sys.getNetwork());
		
		ArrayList<Pair<String,String>> afPairs = sys.getAfPairs();
		for (Pair<String, String> afp : afPairs) {
			principals.add(afp.getFst());
			principals.add(afp.getSnd());
		}
		
		ArrayList<Function> funcs = sys.getFunctions();
		for (Function func : funcs) {
			functionMap.put(func.getName(), func);
		}
		
		ArrayList<Procedure> procs = sys.getProcedures();
		for (Procedure proc : procs) {
			procedureMap.put(proc.getName(), proc);
		}
		
	}
	
	private static void traverseNetwork(Network net)
	{
		Label lb = null;
		if (net instanceof ProcessAtSite) {
			((ProcessAtSite) net).setProcess(simplifyDegenerateParProcess(
					((ProcessAtSite) net).getProcess()));
			traverseProcess(((ProcessAtSite) net).getProcess());
			
			sites.add(((ProcessAtSite) net).getSite().getName());
			principals.add(((ProcessAtSite) net).getSite().getName());
		} else if (net instanceof NewConstN) {
			Constant cst = ((NewConstN) net).getConstant();
			if (cst instanceof ChannelConstant) {
				lb = ((ChannelConstant) cst).getLabel();
				Label simplified = simplifyLabel(lb);
				simplified.setExternal(lb.isExternal());
				((ChannelConstant) cst).setLabel(simplified);
				
				channelMap.put(cst.getName(), (ChannelConstant) cst);
			} else {
				dataMap.put(cst.getName(), (DataConstant) cst);
			}
			Network n = ((NewConstN) net).getNetwork();
			if (n instanceof ParallelNetwork &&
				((ParallelNetwork) n).getComponentNetworks().size() == 1)
			{
				((NewConstN) net).setNetwork(
						((ParallelNetwork) n).getComponentNetworks().get(0));
			}
			traverseNetwork(((NewConstN) net).getNetwork());
		} else {
			assert (net instanceof ParallelNetwork);
			for (Network n : ((ParallelNetwork) net).getComponentNetworks()) {
				traverseNetwork(n);
			}
		}
	}
	
	private static void traverseProcess(Process process)
	{
		Label lb;
		
		if (process instanceof ParallelProcess) {
			
			for (Process p : ((ParallelProcess) process).getProcesses())
				traverseProcess(p);
			
		} else if (process instanceof NewConstP) {
			
			Constant cst = ((NewConstP) process).getConstant();
			if (cst instanceof ChannelConstant) {
				lb = ((ChannelConstant) cst).getLabel();
				Label simplified = simplifyLabel(lb);
				simplified.setExternal(lb.isExternal());
				((ChannelConstant) cst).setLabel(simplified);
				
				channelMap.put(cst.getName(), (ChannelConstant) cst);
			} else {
				dataMap.put(cst.getName(), (DataConstant) cst);
			}
				
			((NewConstP) process).setProcess(simplifyDegenerateParProcess(
					((NewConstP) process).getProcess()));
			traverseProcess(((NewConstP) process).getProcess());
			
		} else if (process instanceof Go) {
			
			((Go) process).setProcess(simplifyDegenerateParProcess(
					((Go) process).getProcess()));
			traverseProcess(((Go) process).getProcess());
			
		} else if (process instanceof Dummy) {
			
			((Dummy) process).setProcess(simplifyDegenerateParProcess(
					((Dummy) process).getProcess()));
			traverseProcess(((Dummy) process).getProcess());
			
		} else if (process instanceof PrefixedProcess) {
			
			traverseBinder(((PrefixedProcess) process).getBinder());
			
			((PrefixedProcess) process).setProcess(
					simplifyDegenerateParProcess(((PrefixedProcess) process).getProcess()));
			traverseProcess(((PrefixedProcess) process).getProcess());
			
		} else if (process instanceof Case) {
			
			traverseExpression(((Case) process).getE());
			
			((Case) process).setProcess1(simplifyDegenerateParProcess(
					((Case) process).getProcess1()));
			((Case) process).setProcess2(simplifyDegenerateParProcess(
					((Case) process).getProcess2()));
			traverseProcess(((Case) process).getProcess1());
			traverseProcess(((Case) process).getProcess2());
			
		} else if (process instanceof ActsFor) {
			
			traverseTerm(((ActsFor) process).getT1());
			traverseTerm(((ActsFor) process).getT2());
			
			((ActsFor) process).setProcess1(simplifyDegenerateParProcess(
					((ActsFor) process).getProcess1()));
			((ActsFor) process).setProcess2(simplifyDegenerateParProcess(
					((ActsFor) process).getProcess2()));
			traverseProcess(((ActsFor) process).getProcess1());
			traverseProcess(((ActsFor) process).getProcess2());
			
		} else if (process instanceof Call) {
			
			for (Expression e : ((Call) process).getArguments()) {
				traverseExpression(e);
			}
			
		} else {
			assert process instanceof Zero;
		}
		
	}
	
	private static void traverseExpression(Expression exp)
	{
		if (exp instanceof Some) {
			traverseTerm(((Some) exp).getTerm());
		} else {
			assert (exp instanceof None || exp instanceof ExpVar);
		}
	}
	
	private static void traverseTerm(Term trm)
	{
		if (trm instanceof FunApp) {
			
			for (Term t: ((FunApp) trm).getArgList()){
				traverseTerm(t);
			}
			
		} else if (trm instanceof ChannelConstant) {
			
			Label original = ((ChannelConstant) trm).getLabel();
			Label simplified = simplifyLabel(original);
			simplified.setExternal(original.isExternal());
			((ChannelConstant) trm).setLabel(simplified);
			
		} else if (trm instanceof Endorse) {

			((Endorse) trm).setL1(simplifyLabel(((Endorse) trm).getL1()));
			((Endorse) trm).setL2(simplifyLabel(((Endorse) trm).getL2()));
			traverseTerm(((Endorse) trm).getT());
			
		} else {
			assert (trm instanceof DataConstant || 
					trm instanceof Principal || 
					trm instanceof TermVar);
		}
	}
	
	private static void traverseBinder(Binder bd)
	{
		if (bd instanceof Out) {
			
			traverseTerm(((Out) bd).getT());
			traverseTerm(((Out) bd).getT1());
			
		} else if (bd instanceof In) {
			
			traverseTerm(((In) bd).getT());
			
		} else if (bd instanceof CompositeBinder) {
			
			((CompositeBinder) bd).setQ(
					simplifyPredicate(((CompositeBinder) bd).getQ()));
			for (Binder b : ((CompositeBinder) bd).getComponentBinders()) {
				traverseBinder(b);
			}
			
		} else {
			
			assert (bd instanceof Assert);
			((Assert) bd).setL1(simplifyLabel(((Assert) bd).getL1()));
			((Assert) bd).setL2(simplifyLabel(((Assert) bd).getL2()));
			traverseBinder(((Assert) bd).getBinder());
	
		}
	}
	
	private static Process simplifyDegenerateParProcess(Process p)
	{
		if (p instanceof ParallelProcess &&
			((ParallelProcess) p).getProcesses().size() == 1)
		{
			return ((ParallelProcess) p).getProcesses().get(0);
		}
		
		return p;
	}
	
	private static QualityPredicate simplifyPredicate(QualityPredicate qp)
	{
		if (qp instanceof NumPredicate) 
			return qp;
		
		QualityPredicate res = qp;
		
		while (true) {
			if (res instanceof AndPredicate &&
				((AndPredicate) res).getOps().size() == 1) 
			{
				res = ((AndPredicate) res).getOps().get(0);
				
			} else if (res instanceof OrPredicate &&
				((OrPredicate) res).getOps().size() == 1) 
			{
				res = ((OrPredicate) res).getOps().get(0);
				
			} else {
				
				ArrayList<QualityPredicate> qps = null;
				if (res instanceof AndPredicate) {
					qps = ((AndPredicate) res).getOps();
					for (int i = 0; i < qps.size(); i++) {
						((AndPredicate) res).setComponentPredicate(
								i, simplifyPredicate(qps.get(i)));
					}
				} else if (res instanceof OrPredicate) {
					qps = ((OrPredicate) res).getOps();
					for (int i = 0; i < qps.size(); i++) {
						((OrPredicate) res).setComponentPredicate(
								i, simplifyPredicate(qps.get(i)));
					}
				}
				
				// if res is a NumPredicate, we can also break here
				break;
			}
			
		}
		
		return res;
	}
	
	public static Label simplifyLabel(Label lb) 
	{
			
		if (lb instanceof AtomicLabel) {
			return lb;
		}
		
		Label res = lb;
		
		while (true) {
			if (res instanceof LabelConjunction &&
				((LabelConjunction) res).getOps().size() == 1) 
			{
				res = ((LabelConjunction) res).getOps().get(0);		
				
			} else if (res instanceof LabelDisjunction &&
				((LabelDisjunction) res).getOps().size() == 1)
			{
				res = ((LabelDisjunction) res).getOps().get(0);
				
			} else {
				ArrayList<Label> lbs = null;
				if (res instanceof LabelConjunction) {
					lbs = ((LabelConjunction) res).getOps();
					for (int i = 0; i < lbs.size(); i++) 
						((LabelConjunction) res).setComponentLabel(
								i, simplifyLabel(lbs.get(i)));	
				} else if (res instanceof LabelDisjunction) {
					lbs = ((LabelDisjunction) res).getOps();
					for (int i = 0; i < lbs.size(); i++) 
						((LabelDisjunction) res).setComponentLabel(
								i, simplifyLabel(lbs.get(i)));
				} 
				
				// if res is an AtomicLabel, we can also break here
				break;
			}
		}
		
		return res;
	}

}
